import os
import logging
from PIL import Image, ImageOps
import numpy as np
import boto3

from label_studio_ml.model import LabelStudioMLBase
from label_studio_ml.response import ModelResponse

from control_models.base import ControlModel
from control_models.choices import ChoicesModel
from control_models.rectangle_labels import RectangleLabelsModel
from control_models.rectangle_labels_obb import RectangleLabelsObbModel
from control_models.polygon_labels import PolygonLabelsModel
from control_models.keypoint_labels import KeypointLabelsModel
from control_models.keypoint_polygon_labels import KeypointPolygonLabelsModel
from control_models.video_rectangle import VideoRectangleModel
from control_models.timeline_labels import TimelineLabelsModel
from control_models.paddleocr_labels import PaddleOCRModel
from typing import List, Dict, Optional

LABEL_STUDIO_ACCESS_TOKEN = os.environ.get("LABEL_STUDIO_ACCESS_TOKEN")
LABEL_STUDIO_HOST = os.environ.get("LABEL_STUDIO_HOST")

logger = logging.getLogger()
if not os.getenv("LOG_LEVEL"):
    logger.setLevel(logging.INFO)

# Register available model classes
available_model_classes = [
    ChoicesModel,
    RectangleLabelsModel,
    RectangleLabelsObbModel,
    PolygonLabelsModel,
    KeypointLabelsModel,
    KeypointPolygonLabelsModel,
    VideoRectangleModel,
    TimelineLabelsModel,
]

AWS_ACCESS_KEY_ID = os.environ.get("AWS_ACCESS_KEY_ID")
AWS_SECRET_ACCESS_KEY = os.environ.get("AWS_SECRET_ACCESS_KEY")
AWS_SESSION_TOKEN = os.environ.get("AWS_SESSION_TOKEN")
AWS_ENDPOINT = os.environ.get("AWS_ENDPOINT")

# S3_TARGET = boto3.resource('s3',
#                            endpoint_url=AWS_ENDPOINT,
#                            aws_access_key_id=AWS_ACCESS_KEY_ID,
#                            aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
#                            aws_session_token=AWS_SESSION_TOKEN,
#                            config=boto3.session.Config(signature_version='s3v4'),
#                            verify=False)


class YOLO(LabelStudioMLBase):
    """Label Studio ML Backend based on Ultralytics YOLO"""
    MODEL_DIR = os.environ.get('MODEL_DIR', '.')
    def __init__(self, **kwargs):
        super().__init__(** kwargs)
        # 配置Label Studio访问令牌
        self.access_token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ0b2tlbl90eXBlIjoicmVmcmVzaCIsImV4cCI6ODA2MjM0MDM2MSwiaWF0IjoxNzU1MTQwMzYxLCJqdGkiOiI5YmMxZTQ3Y2M0ZTY0ZDEzOGQ2NjUxMDIxYjgwOWEzOCIsInVzZXJfaWQiOjF9.HWO8LBON5pQp4CqJhBxEpwKAKJ5W5zcd1no65Jw0LYI"  # 替换为实际令牌
        # 配置Label Studio服务器地址
        # self.label_studio_host = "http://192.188.1.196:8080"  # 替换为实际地址


    def setup(self):
        """Configure any parameters of your model here"""
        self.set("model_version", "yolo")

    def detect_control_models(self,models = "yolo") -> List[ControlModel]:
        """Detect control models based on the labeling config.
        Control models are used to predict regions for different control tags in the labeling config.
        """

        if models == "paddleocr":
            control_models = []

            for control in self.label_interface.controls:
                # skipping tags without toName
                instance = PaddleOCRModel.create(self, control)
                control_models.append(instance)
            return control_models
        elif models == "yolo":
            control_models = []

            for control in self.label_interface.controls:
                # skipping tags without toName
                if not control.to_name:
                    logger.warning(
                        f'{control.tag} {control.name} has no "toName" attribute, skipping it'
                    )
                    continue

                # match control tag with available control models
                for model_class in available_model_classes:
                    mode = self.extra_params.get("mode", None)
                    if model_class.is_control_matched(control, mode) and model_class.mode == mode :
                        instance = model_class.create(self, control)
                        if not instance:
                            logger.debug(
                                f"No instance created for {control.tag} {control.name}"
                            )
                            continue
                        if not instance.label_map:
                            logger.error(
                                f"No label map built for the '{control.tag}' control tag '{instance.from_name}'.\n"
                                f"This indicates that your Label Studio config labels do not match the model's labels.\n"
                                f"To fix this, ensure that the 'value' or 'predicted_values' attribute "
                                f"in your Label Studio config matches one or more of these model labels.\n"
                                f"If you don't want to use this control tag for predictions, "
                                f'add `model_skip="true"` to it.\n'
                                f"Examples:\n"
                                f'  <Label value="Car"/>\n'
                                f'  <Label value="YourLabel" predicted_values="label1,label2"/>\n'
                                f"Labels provided in your labeling config:\n"
                                f"  {str(control.labels_attrs)}\n"
                                f"Available '{instance.model_path}' model labels:\n"
                                f"  {list(instance.model.names.values())}"
                            )
                            continue

                        control_models.append(instance)
                        logger.debug(f"Control tag with model detected: {instance}")
                        break

            if not control_models:
                control_tags = ", ".join([c.type for c in available_model_classes])
                raise ValueError(
                    f"No suitable control tags (e.g. {control_tags} connected to Image or Video object tags) "
                    f"detected in the label config:\n{self.label_config}"
                )

            return control_models

    def predict(
        self, tasks: List[Dict], context: Optional[Dict] = None, **kwargs
    ) -> ModelResponse:
        """Run YOLO predictions on the tasks
        :param tasks: [Label Studio tasks in JSON format](https://labelstud.io/guide/task_format.html)
        :param context: [Label Studio context in JSON format](https://labelstud.io/guide/ml_create)
        :return model_response
            ModelResponse(predictions=predictions) with
            predictions [Predictions array in JSON format]
            (https://labelstud.io/guide/export.html#Label-Studio-JSON-format-of-annotated-tasks)
        """
        
        # context = kwargs.get('context')
        if context: # 如果有上下文就标注
            from_name, to_name, value = self.label_interface.get_first_tag_occurence(
                'TextArea',
                'Image'
            )
            if context["result"]:
                # return []
                control_models:list[PaddleOCRModel] = self.detect_control_models("paddleocr")
                for model in control_models:
                    task = tasks[0]
                    img_path_url = task["data"][value]
                    image = self.load_image(img_path_url, task.get('id'))

                    result = context.get('result')[-1]
                    meta = self._extract_meta({**task, **result})
                    x = meta["x"] * meta["original_width"] / 100
                    y = meta["y"] * meta["original_height"] / 100
                    w = meta["width"] * meta["original_width"] / 100
                    h = meta["height"] * meta["original_height"] / 100

                    result_text = model.predict_ocr(
                        np.array(image.crop((x, y, x + w, y + h)), dtype=np.uint8),
                    )
                    meta["text"] = result_text["rec_texts"]
                    score = result_text["rec_scores"]

                    temp = {
                        "original_width": meta["original_width"],
                        "original_height": meta["original_height"],
                        "image_rotation": 0,
                        "value": {
                            "points": result["value"]["points"],
                            "closed": result["value"]["closed"],
                            "rotation": 0,
                            "text": [
                                meta["text"]
                            ]
                        },
                        "id": meta["id"],
                        "from_name": from_name,
                        "to_name": meta['to_name'],
                        "type": "textarea",
                        "origin": "manual"
                    }
                    return [{
                        'result': [temp, result],
                        'score': score,
                        'model_version': self.get('model_version')
                    }]
                
            
        logger.info(
            f"Run prediction on {len(tasks)} tasks, project ID = {self.project_id}"
        )
        control_models = self.detect_control_models()
        
        predictions = []
        for task in tasks:

            regions = []
            for model in control_models:
                path = model.get_path(task)
                regions += model.predict_regions(path)

            # calculate final score
            all_scores = [region["score"] for region in regions if "score" in region]
            avg_score = sum(all_scores) / max(len(all_scores), 1)

            # compose final prediction
            prediction = {
                "result": regions,
                "score": avg_score,
                "model_version": self.model_version,
            }
            predictions.append(prediction)

        return ModelResponse(predictions=predictions)

    def load_image(self, img_path_url, task_id):
        # load an s3 image, this is very basic demonstration code
        # you may need to modify to fit your own needs
        if img_path_url.startswith("s3:"):
            bucket_name = img_path_url.split("/")[2]
            key = "/".join(img_path_url.split("/")[3:])

            obj = S3_TARGET.Object(bucket_name, key).get()
            data = obj['Body'].read()
            image = Image.open(io.BytesIO(data))
            image = ImageOps.exif_transpose(image)
            return image
        else:
            cache_dir = os.path.join(self.MODEL_DIR, '.file-cache')
            os.makedirs(cache_dir, exist_ok=True)
            logger.debug(f'Using cache dir: {cache_dir}')
            filepath = self.get_local_path(
                img_path_url,
                cache_dir=cache_dir,
                ls_access_token=LABEL_STUDIO_ACCESS_TOKEN,
                ls_host=LABEL_STUDIO_HOST,
                task_id=task_id
            )
            image = Image.open(filepath)
            image = ImageOps.exif_transpose(image)
            return image

    @staticmethod
    def _extract_meta(task):
        meta = dict()
        if task:
            logger.info(f"task: {task}")

            meta['id'] = task['id']
            meta['from_name'] = task['from_name']
            meta['type'] = task['type']
            meta['to_name'] = task['to_name']
            points = task['value']["points"]
            meta['x'] = min([p[0] for p in points])
            meta['y'] = min([p[1] for p in points])
            meta['width'] = max([p[0] for p in points]) - meta['x']
            meta['height'] = max([p[1] for p in points]) - meta['y']
            meta["original_width"] = task['original_width']
            meta["original_height"] = task['original_height']
        return meta

    def fit(self, event, data, **kwargs):
        """
        This method is called each time an annotation is created or updated.
        Or it's called when "Start training" clicked on the model in the project settings.
        """
        results = {}
        control_models = self.detect_control_models()
        for model in control_models:
            training_result = model.fit(event, data, **kwargs)
            results[model.from_name] = training_result

        return results
